from autils.v8 import v8

import cv2

import threading
import pyttsx4
import time
import os


class Predict(object):
    """docstring for Predict"""
    def __init__(self, model_path, source="0"):
        #super(Predict, self).__init__()
        #self.arg = arg
        # self.class_names_list = ["bucket", "hook", "person", "cement_mixer", "excavator", "crane"]
        self.YOLO = v8(model_path)
        self.model = self.YOLO.model

        self.class_names_list = self.model.names
        self.source = source
        self.cap = cv2.VideoCapture(6)

        self.danger_distance = []
        self.old_path = ""
        self.text = ""
        self.danger_count_status = True
        self.picture_style = [".png", ".jpg"]
        self.x_sta = ""

        self.start_time = time.time()
        self.end_time = time.time()

        self.time_lock = False

    def choose_video_predict(self, cvmode):

        if cvmode == self.x_sta:
            return "Error"
        
        self.video = cv2.VideoCapture(self.source) # 视频
        self.old_path = self.source

        self.x_sta = cvmode # remeber to adjust

    def change_model(self, model_path):
        self.YOLO = v8(model_path)
        self.model = self.YOLO.model
        self.class_names_list = self.model.names
        

    def frame(self):

        if self.source == "0":
            # print("================camera===================")
            _, self.color_image = self.cap.read()

        elif os.path.splitext(str(self.source))[-1] in self.picture_style:
            # pass
            self.old_path = self.source # 保证旧的文件得到更新
            self.color_image = cv2.imread(self.source)

        elif os.path.splitext(str(self.source))[-1] == ".mp4":
            if self.old_path != self.source:
                self.x_sta = ""
            mode_count_status = "2"
            self.choose_video_predict(mode_count_status)
            _, self.color_image = self.video.read()
            
        # 创建传递图像
        self.img = self.color_image.copy()

    def det_data(self):
        self.detect_end_status = False
        self.frm = None
        self.label_array = []
        self.person_xyz_array = []

        '''person'''
        self.person_message_points = []
        self.del_index_list = []

    def chinese_l(self, text):
        en_l = ["bucket", "hook", "person", "cement_mixer", "excavator", "crane"]
        chinese_l = ["铲子", "吊钩", "人", "水泥车", "挖掘机", "吊车"]
        index = en_l.index(text)
        return chinese_l[index]

    def predict(self, self_source, self_confidence, self_r, self_g, self_b):
        self.source = self_source
        self.det_data() # refrash data
        self.frame() # refrash frame
        metrics = self.model.predict(self.img)

        for m in metrics:
            positions_box = m.boxes.cpu()
            class_name = positions_box.cls.tolist()
            conf = positions_box.conf.tolist()
            xyxy = positions_box.xyxy.tolist()
            for index,x in enumerate(xyxy):
                if float(conf[index]) > self_confidence:
                    # position
                    a,b,c,d = int(x[0]), int(x[1]), int(x[2]), int(x[3])
                    # data_info
                    label_list = [self.class_names_list[int(class_name[index])], round(float(conf[index]),2),[a,b,c,d]]
                    self.label_array.append(label_list)
                    c1, c2 = (a,b), (c,d)
            # print(class_name)
            # print(xyxy)
        self.frm = self.img
        # except
        if self.img is None or self.img.size == 0:
            print("Error: Frame is empty!")
            return None, self.color_image

        self.end_time = time.time()

        # define color

        rag_color = (self_b*255,self_g*255,self_r*255)
        '''end'''

        '''reg circle'''
        
        for index,i in enumerate(self.label_array):
            x1,y1,x2,y2 = i[-1]
            # cv2.rectangle(self.frm,(x1,y1),(x2,y2),rag2_color,2)
            # cv2.putText(self.frm,o[0],(x1,y1 - 10),cv2.FONT_HERSHEY_SIMPLEX,0.8,rag2_color,2)
            cv2.rectangle(self.frm,(x1,y1),(x2,y2),rag_color,2)
            cv2.putText(self.frm,i[0],(x1,y1 - 10),cv2.FONT_HERSHEY_SIMPLEX,0.8,rag_color,2)
        
        # print((self.end_time-self.start_time))
        if (self.end_time-self.start_time) % 2 > 0.0 and (self.end_time-self.start_time) % 2 < 0.1:
            text = ""
            for i in self.label_array:
                text += f"{i[0]}: {i[1]}\n\n"
            self.text = text + "Done"
        
        # 临时
        # ['crane', 0.162692962355475, [365, 282, 619, 423], True]
        # cv2.rectangle(self.frm,(365,282),(619,423),(0,255,0),2)
        # cv2.putText(self.frm,'crane',(365,282 - 10),cv2.FONT_HERSHEY_SIMPLEX,0.8,(0,255,0),2)

        print("\n\n")

        '''reg end'''

        '''text'''

        # if self.danger_count_status:
            # print("危险！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！")
        self.danger_count_status = False
        temp_list = []

        self.danger_distance = temp_list
        # print(self.danger_distance)
        # print(f'detect: {self.label_array}')
        return self.frm, self.color_image, self.label_array
